use crate::ast::*;
use crate::error::{CompilerError, Result};
use super::TypeInference;
use super::super::TypeEnvironment;

impl TypeInference {


    /// Infer types in a block
    /// If expected_return_type is provided, use it for context-aware return type inference
    pub(super) fn infer_block(&mut self, block: &Block, env: &mut TypeEnvironment) -> Result<Type> {
        self.infer_block_with_expected(block, env, None)
    }


    /// Infer types in a block with expected return type
    pub(super) fn infer_block_with_expected(
        &mut self,
        block: &Block,
        env: &mut TypeEnvironment,
        expected_return_type: Option<&Type>,
    ) -> Result<Type> {
        let mut last_type = Type::Basic(crate::ast::types::BasicType::Void);

        for statement in &block.statements {
            // For return statements, use expected return type for context-aware inference
            if let crate::ast::stmt::Statement::Return(return_stmt) = statement {
                if let Some(value) = &return_stmt.value {
                    let inferred =
                        self.infer_expression_with_expected(value, env, expected_return_type)?;
                    last_type = inferred;
                    continue;
                }
            }
            last_type = self.infer_statement(statement, env)?;
        }

        Ok(last_type)
    }


    /// Infer types in a statement
    pub(super) fn infer_statement(
        &mut self,
        statement: &crate::ast::stmt::Statement,
        env: &mut TypeEnvironment,
    ) -> Result<Type> {
        match statement {
            crate::ast::stmt::Statement::Expression(expr) => self.infer_expression(expr, env),
            crate::ast::stmt::Statement::VariableDecl(var) => {
                self.infer_variable_declaration(var, env)
            }
            crate::ast::stmt::Statement::ConstantDecl(const_) => {
                self.infer_constant_declaration(const_, env)
            }
            crate::ast::stmt::Statement::Return(return_stmt) => {
                self.infer_return_statement(return_stmt, env)
            }
            crate::ast::stmt::Statement::Block(block) => {
                let mut block_env = env.child();
                self.infer_block_statement(block, &mut block_env)
            }
            _ => {
                // Other statements return void
                Ok(Type::Basic(crate::ast::types::BasicType::Void))
            }
        }
    }


    /// Infer types in a variable declaration statement
    pub(super) fn infer_variable_declaration(
        &mut self,
        var: &crate::ast::stmt::VariableDeclStmt,
        env: &mut TypeEnvironment,
    ) -> Result<Type> {
        let var_type = if let Some(declared_type) = &var.var_type {
            // If variable has type annotation, use it for context-aware type inference
            let declared_type_clone = declared_type.clone();
            if let Some(initializer) = &var.initializer {
                // Infer initializer with expected type (for context-aware Result type inference)
                let _inferred_type = self.infer_expression_with_expected(
                    initializer,
                    env,
                    Some(&declared_type_clone),
                )?;

                // Check that inferred type matches declared type
                // TODO: Add type compatibility check here
                declared_type_clone
            } else {
                declared_type_clone
            }
        } else if let Some(initializer) = &var.initializer {
            self.infer_expression(initializer, env)?
        } else {
            return Err(CompilerError::type_error(
                var.location.line,
                var.location.column,
                "Variable declaration must have type annotation or initializer",
            ));
        };

        env.bind_variable(var.name.clone(), var_type.clone());
        Ok(Type::Basic(crate::ast::types::BasicType::Void))
    }


    /// Infer types in a constant declaration statement
    pub(super) fn infer_constant_declaration(
        &mut self,
        const_: &crate::ast::stmt::ConstantDeclStmt,
        env: &mut TypeEnvironment,
    ) -> Result<Type> {
        let const_type = if let Some(declared_type) = &const_.const_type {
            declared_type.clone()
        } else {
            self.infer_expression(&const_.value, env)?
        };

        env.bind_variable(const_.name.clone(), const_type);
        Ok(Type::Basic(crate::ast::types::BasicType::Void))
    }


    /// Infer types in a return statement
    /// Uses function return type from environment for context-aware type inference
    pub(super) fn infer_return_statement(
        &mut self,
        return_stmt: &crate::ast::stmt::ReturnStmt,
        env: &mut TypeEnvironment,
    ) -> Result<Type> {
        if let Some(value) = &return_stmt.value {
            // Note: Return statement context is handled in infer_block_with_expected
            // This method is called from infer_statement, but context-aware inference
            // happens in infer_block_with_expected when processing return statements
            self.infer_expression(value, env)
        } else {
            Ok(Type::Basic(crate::ast::types::BasicType::Void))
        }
    }


    /// Infer types in a block statement
    pub(super) fn infer_block_statement(
        &mut self,
        block: &crate::ast::stmt::BlockStmt,
        env: &mut TypeEnvironment,
    ) -> Result<Type> {
        let mut last_type = Type::Basic(crate::ast::types::BasicType::Void);

        for statement in &block.statements {
            last_type = self.infer_statement(statement, env)?;
        }

        Ok(last_type)
    }

}
